#include "sm2.h"

#include <assert.h>
#include <memory.h>

#include "sm3.h"

static void MessageHash(const uint8_t ZA[SM3_HASH_LENGTH], const uint8_t *m, size_t mlen, uint8_t e[SM3_HASH_LENGTH]) {
	SM3_WorkArea sm3;
	SM3_Initialize(&sm3);


	SM3_Input(&sm3, ZA, SM3_HASH_LENGTH);
	SM3_Input(&sm3, m, mlen);
	SM3_Result(&sm3, e);
}


void SM2_Initialize(SM2Work *sm2, size_t size, const uint8_t *G, uint8_t* n, const ECWork *ec) {
	assert(sm2 != NULL);

	sm2->size = (size+7)/8;

	UInteger_FromBytes(&(sm2->G.x), G, EC_SIZE(ec));
	UInteger_FromBytes(&(sm2->G.y), G+EC_SIZE(ec), EC_SIZE(ec));
	UInteger_FromBytes(&(sm2->n), n, sm2->size);

	sm2->ec = ec;

}

void SM2_ZA(SM2Work *sm2, const uint8_t*  id, size_t idlen, const uint8_t *PA, uint8_t ZA[SM3_HASH_LENGTH]) {

	size_t tstrsize = EC_SIZE(sm2->ec) > sm2->size ? EC_SIZE(sm2->ec) : sm2->size;
	uint8_t tstr[tstrsize];

	SM3_WorkArea sm3;
	SM3_Initialize(&sm3);

	uint16_t entlen = idlen*8;
	tstr[0] = entlen >> 8;
	tstr[1] = entlen & 0xFF;
	SM3_Input(&sm3, tstr, 2);

	SM3_Input(&sm3, id, idlen);
	
	size_t n = sm2->ec->za(sm2->ec, NULL, 0);
	uint8_t ecza[n];
	sm2->ec->za(sm2->ec, ecza, n);
	SM3_Input(&sm3, ecza, n);

	UInteger_ToBytes(&(sm2->G.x), tstr, EC_SIZE(sm2->ec));
	SM3_Input(&sm3, tstr, EC_SIZE(sm2->ec));
	
	UInteger_ToBytes(&(sm2->G.y), tstr, EC_SIZE(sm2->ec));
	SM3_Input(&sm3, tstr, EC_SIZE(sm2->ec));

	SM3_Input(&sm3, PA, EC_SIZE(sm2->ec)*2);
	
	SM3_Result(&sm3, ZA);
}

void SM2_PublicKey(SM2Work *sm2, const uint8_t *dA, uint8_t *PA) {
	 
		UInteger dAi;
		UInteger_FromBytes(&dAi, dA, sm2->size);

		ECPoint PApt;
		EC_MUL(sm2->ec, &dAi, &(sm2->G), &PApt);

		UInteger_ToBytes(&(PApt.x), PA, EC_SIZE(sm2->ec));
		UInteger_ToBytes(&(PApt.y), PA+EC_SIZE(sm2->ec), EC_SIZE(sm2->ec));

}

void SM2_Sign(const SM2Work* sm2, const uint8_t *dA, const uint8_t ZA[SM3_HASH_LENGTH], const uint8_t *m, size_t mlen, const uint8_t *k, uint8_t *rs) {
	uint8_t tstr[sm2->size];

	uint8_t e[SM3_HASH_LENGTH];

	MessageHash(ZA, m, mlen, e);

	UInteger ki;
	ECPoint pt;

	UInteger_FromBytes(&ki, k, sm2->size);
	EC_MUL(sm2->ec, &ki, &(sm2->G), &pt);

	UInteger ri, si, ei;
	UInteger_FromBytes(&ei, e, sm2->size);
	UInteger_ModAdd(&ei, &(pt.x), &(sm2->n), &ri);
	
	// ((k-r*dA)/(1+dA) mod n)
	UInteger one;
	UInteger_FromUInt(&one, 1);
	UInteger dAi;
	UInteger_FromBytes(&dAi, dA, sm2->size);
	UInteger t[2];
	UInteger_Add(&one, &dAi, t);
	UInteger_ModMul(&ri, &dAi, &(sm2->n), t+1);
	UInteger_ModSub(&ki, t+1, &(sm2->n), t+1);
	UInteger_ModDiv(t+1, t, &(sm2->n), &si);

	UInteger_ToBytes(&ri, rs, sm2->size);
	UInteger_ToBytes(&si, rs+sm2->size, sm2->size);


}



bool SM2_Verify(const SM2Work* sm2, const uint8_t *PA, const uint8_t ZA[SM3_HASH_LENGTH], const uint8_t *m, size_t mlen, const uint8_t *rs) {

	uint8_t e[SM3_HASH_LENGTH];

	MessageHash(ZA, m, mlen, e);

	UInteger ei;
	UInteger_FromBytes(&ei, e, sm2->size);

	UInteger ri, si, ti;
	UInteger_FromBytes(&ri, rs, sm2->size);
	UInteger_FromBytes(&si, rs+sm2->size, sm2->size);
	UInteger_ModAdd(&ri, &si, &(sm2->n), &ti);

	ECPoint PApt;
	UInteger_FromBytes(&(PApt.x), PA, EC_SIZE(sm2->ec));
	UInteger_FromBytes(&(PApt.y), PA+EC_SIZE(sm2->ec), EC_SIZE(sm2->ec));

	ECPoint pt[2];
	EC_MUL(sm2->ec, &si, &(sm2->G), pt);
	EC_MUL(sm2->ec, &ti, &PApt, pt+1);
	EC_ADD(sm2->ec, pt, pt+1, pt);

	UInteger R;
	UInteger_ModAdd(&ei, &(pt[0].x), &(sm2->n), &R);

	return UInteger_Compare(&ri, &R) == 0;
	
}


/**
 * 密钥派生函数，参考第4部分《公钥加密算法》5.4.3节
 */
void SM2_KDF(const uint8_t *z, size_t zsize, uint8_t *k, size_t ksize) {

	size_t number = (ksize+SM3_HASH_LENGTH-1)/SM3_HASH_LENGTH;

	for (uint32_t ct=1; ct<=number; ++ct) {
		uint8_t ct8s[4] = { ct>>24, ct>>16, ct>>8, ct };

		uint8_t h[SM3_HASH_LENGTH];

		SM3_WorkArea sm3;
		SM3_Initialize(&sm3);
		SM3_Input(&sm3, z, zsize);
		SM3_Input(&sm3, ct8s, 4);
		SM3_Result(&sm3, h);

		size_t offset = SM3_HASH_LENGTH*(ct-1);
		size_t len = ksize >= offset+SM3_HASH_LENGTH ? SM3_HASH_LENGTH : ksize-offset;
		memcpy(k+offset, h, len); 
	}

}

size_t SM2_CipherStubLength(const SM2Work *sm2) {
	return 1 + 2*EC_SIZE(sm2->ec) + SM3_HASH_LENGTH;
}

bool SM2_Encrypt(const SM2Work *sm2, const uint8_t *publicKey, const uint8_t *random, uint8_t *plain, size_t plainLength, uint8_t* cipher) {
	uint8_t *C1 = cipher;
	uint8_t *C2 = cipher + EC_SIZE(sm2->ec)*2 + 1;
	uint8_t *C3 = C2 + plainLength;
	UInteger k;

	UInteger_FromBytes(&k, random,  sm2->size);

	// A.1 检查 1 <= k <= n-1

	// A.2 计算 C1 = k*G
	ECPoint C1i;
	EC_MUL(sm2->ec, &k, &(sm2->G), &C1i);
	C1[0] = 0x04;
	UInteger_ToBytes(&(C1i.x), C1+1, EC_SIZE(sm2->ec));
	UInteger_ToBytes(&(C1i.y), C1+1+EC_SIZE(sm2->ec), EC_SIZE(sm2->ec));

	// A.3 检查S = h*PB, S是否是无穷远点

	// A4 计算(x2,y2) = k*PB
	ECPoint PB, XY2;
	UInteger_FromBytes(&(PB.x), publicKey, EC_SIZE(sm2->ec));
	UInteger_FromBytes(&(PB.y), publicKey+EC_SIZE(sm2->ec), EC_SIZE(sm2->ec));
	EC_MUL(sm2->ec, &k, &PB, &XY2);

	uint8_t xystr[EC_SIZE(sm2->ec)*2];
	UInteger_ToBytes(&(XY2.x), xystr, EC_SIZE(sm2->ec));
	UInteger_ToBytes(&(XY2.y), xystr+EC_SIZE(sm2->ec), EC_SIZE(sm2->ec));
	SM2_KDF(xystr, EC_SIZE(sm2->ec)*2, C2, plainLength);

	// A6
	for (size_t i=0; i<plainLength; ++i) {
		C2[i] ^= plain[i];
	}

	// A7
	SM3_WorkArea sm3;
	SM3_Initialize(&sm3);
	SM3_Input(&sm3, xystr, EC_SIZE(sm2->ec));
	SM3_Input(&sm3, plain, plainLength);
	SM3_Input(&sm3, xystr+EC_SIZE(sm2->ec), EC_SIZE(sm2->ec));
	SM3_Result(&sm3, C3);


	return true;
}

bool SM2_Decrypt(const SM2Work *sm2, const uint8_t *privateKey, const uint8_t *cipher, size_t cipherLength, uint8_t *plain) {

	size_t plainLength = cipherLength - SM2_CipherStubLength(sm2);
	const uint8_t *C1 = cipher;
	const uint8_t *C2 = C1 + 1 + 2*EC_SIZE(sm2->ec);
	const uint8_t *C3 = C2 + plainLength;

	// B1 判断C1是否满足曲线
	ECPoint C1i;
	UInteger_FromBytes(&(C1i.x), C1+1, EC_SIZE(sm2->ec));
	UInteger_FromBytes(&(C1i.y), C1+1+EC_SIZE(sm2->ec), EC_SIZE(sm2->ec));

	EC_ON(sm2->ec, &C1i);

	// B2

	// B3
	UInteger dB;
	UInteger_FromBytes(&(dB), privateKey, sm2->size);
	ECPoint XY2;
	EC_MUL(sm2->ec, &dB, &C1i, &XY2);

	uint8_t xy2str[sm2->size*2];
	UInteger_ToBytes(&(XY2.x), xy2str, EC_SIZE(sm2->ec));
	UInteger_ToBytes(&(XY2.y), xy2str+EC_SIZE(sm2->ec), EC_SIZE(sm2->ec));

	// B4
	SM2_KDF(xy2str, EC_SIZE(sm2->ec)*2, plain, plainLength);

	// B5
	for (size_t i=0; i<plainLength; ++i) {
		plain[i] ^= C2[i];
	}

	// B6
	uint8_t h[SM3_HASH_LENGTH];
	SM3_WorkArea sm3;
	SM3_Initialize(&sm3);
	SM3_Input(&sm3, xy2str, EC_SIZE(sm2->ec));
	SM3_Input(&sm3, plain, plainLength);
	SM3_Input(&sm3, xy2str+EC_SIZE(sm2->ec), EC_SIZE(sm2->ec));
	SM3_Result(&sm3, h);

	return memcmp(h, C3, SM3_HASH_LENGTH) == 0;
}
